import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import logging
from datetime import datetime
import pinecone
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from pinecone import Pinecone, ServerlessSpec
from collections import Counter

# 配置日志记录器
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

pinecone = Pinecone(api_key="acd1bdc5-a876-42db-8b49-191bd4bb1065")

# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
# 这个if是否需要，要看情况而定
# 比如有的时候，如果不想要重复删除在创建，这个if就可以不要
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 创建索引实例
index = pinecone.Index(index_name)

# 向索引中添加训练数据
vectors = []
for i, (vec, label) in enumerate(zip(X_train, y_train)):
    # 确保label是Python内置的int类型
    vectors.append((str(i), vec.tolist(), {"label": int(label)}))
index.upsert(vectors)


def predict_digit(stroke):
    """
    接收一个手写数字的笔画数据，并使用Pinecone进行预测。
    参数：
        stroke (list): 包含用户手写笔画数据的列表。
    返回：
        str: 预测的数字。
    """
    # 创建一个空白图像
    image = Image.new('RGB', (280, 280), color=(255, 255, 255))
    draw = ImageDraw.Draw(image)

    # 绘制每一笔画
    for points in stroke:
        draw.line(points, fill="black", width=5)

    # 将图像转换为灰度模式并调整大小为8x8像素
    image = image.resize((8, 8), Image.Resampling.LANCZOS).convert('L')
    # 将图像转换为numpy数组并进行归一化
    image_array = np.array(image) / 255.0
    # 将图像数组转换为二维数组
    image_array = image_array.flatten()

    # 使用Pinecone进行预测
    results = index.query(image_array.tolist(), top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    final_prediction = Counter(labels).most_common(1)[0][0]
    return str(final_prediction)

# 创建Gradio界面
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="Draw a digit"),
    outputs="text",
    title="Handwritten Digit Recognition",
    description="Draw a digit in the box and we'll tell you what it is."
)

# 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch()
